import glob 
import numpy as np 
from keras.utils import to_categorical
from typing import Tuple, List
from ..utils import R

class DataGenerator:
    def __init__(self, data_dir, samples) -> None:
        self.data_dir = data_dir
        self.samples = samples
        self.filenames = set(filename for filename, index in samples)
        self.num_samples = None 

    def GetNumberSamples(self, batch_size = 128, num_classes=19*19):
        if self.num_samples is not None:
            return self.num_samples
        else:
            self.num_samples = 0
            for x, y in self._Generate(batch_size, num_classes):
                self.num_samples += x.shape[0]
            return self.num_samples
        
    def _Generate(self, batch_size, num_classes):
        for zip_filename in self.filenames:
            filename = zip_filename.replace('.tar.gz', '') + 'train'
            base = self.data_dir + '/' + filename + '_features_*.npy'
            for feature_file in glob.glob(base):
                label_file = feature_file.replace('features', 'labels')
                x = np.load(feature_file)
                y = np.load(label_file)
                x = x.astype('float32')
                y = to_categorical(y.astype(int), num_classes)
                while x.shape[0] >= batch_size:
                    x_batch, x = x[:batch_size], x[batch_size:]
                    y_batch, y = y[:batch_size], y[batch_size:]
                    yield x_batch, y_batch

    def Generate(self, batch_size=128, num_classes=19*19):
        while True:
            for item in self._Generate(batch_size, num_classes):
                yield item 